import pandas as pd
import numpy as np
import akshare as ak
import matplotlib.pyplot as plt
import talib
from datetime import datetime, timedelta
import warnings
from sklearn.linear_model import LinearRegression
warnings.filterwarnings('ignore')

plt.rcParams['font.family'] = 'DejaVu Sans'
plt.rcParams['figure.figsize'] = (12, 8)

class EnhancedCSI300Strategy:
    def __init__(self, symbol="sh000300", start_date="20050601"):
        self.symbol = symbol
        self.start_date = start_date
        self.data = self.download_data()
        self.calculate_indicators()
        self.identify_market_regime()
        self.data.to_csv('data.csv', index=False, encoding='utf-8-sig') 
        
    def download_data(self):
        """下载数据"""
        df = ak.stock_zh_index_daily(symbol=self.symbol)
        df = df.sort_index(ascending=True)
        
        # 确保有date列并设置为索引
        if 'date' not in df.columns:
            df['date'] = df.index
        df['date'] = pd.to_datetime(df['date'])
        df['date_col'] = df['date']
        df.set_index('date', inplace=True)
        cols = df.columns.tolist()
        cols.remove('date_col')
        cols.insert(0, 'date_col')
        df = df[cols]
        
        # 筛选起始日期
        df = df[df.index >= self.start_date]
        
        df = df.rename(columns={
            'close': 'Close',
            'open': 'Open',
            'high': 'High',
            'low': 'Low',
            'volume': 'Volume'
        })
        
        # 确保索引是日期时间类型
        df.index = pd.to_datetime(df.index)
        return df
    
    def calculate_indicators(self):
        """计算技术指标"""
        # 移动平均线
        self.data['MA20'] = talib.MA(self.data['Close'], timeperiod=20)
        self.data['MA50'] = talib.MA(self.data['Close'], timeperiod=50)
        self.data['MA120'] = talib.MA(self.data['Close'], timeperiod=120)
        self.data['MA5'] = talib.MA(self.data['Close'], timeperiod=5)
        
        # MACD
        self.data['MACD'], self.data['MACD_Signal'], self.data['MACD_Hist'] = talib.MACD(
            self.data['Close'], fastperiod=12, slowperiod=26, signalperiod=9)
        
        # RSI
        self.data['RSI'] = talib.RSI(self.data['Close'], timeperiod=14)
        
        # 布林带
        self.data['BB_Upper'], self.data['BB_Middle'], self.data['BB_Lower'] = talib.BBANDS(
            self.data['Close'], timeperiod=20, nbdevup=2, nbdevdn=2, matype=0)
        
        # 成交量指标
        self.data['Volume_MA20'] = talib.MA(self.data['Volume'], timeperiod=20)
        self.data['Volume_Ratio'] = self.data['Volume'] / self.data['Volume_MA20']
        
        # CCI
        self.data['CCI'] = talib.CCI(self.data['High'], self.data['Low'], self.data['Close'], timeperiod=20)
        
        # ADX - 趋势强度
        self.data['ADX'] = talib.ADX(self.data['High'], self.data['Low'], self.data['Close'], timeperiod=14)
        
        # ATR - 波动率指标
        self.data['ATR'] = talib.ATR(self.data['High'], self.data['Low'], self.data['Close'], timeperiod=14)
        self.data['ATR_Pct'] = self.data['ATR'] / self.data['Close']
        
        # 价格位置相对于布林带
        self.data['BB_Position'] = (self.data['Close'] - self.data['BB_Lower']) / (
            self.data['BB_Upper'] - self.data['BB_Lower'])
        
        # 动量指标
        self.data['Momentum'] = self.data['Close'] / self.data['Close'].shift(5) - 1
        
        # 市场情绪指标 - 使用价格与均线的偏离度
        self.data['Price_MA20_Ratio'] = self.data['Close'] / self.data['MA20'] - 1
        self.data['Price_MA50_Ratio'] = self.data['Close'] / self.data['MA50'] - 1
        
        # 资金流指标 - 使用OBV
        self.data['OBV'] = talib.OBV(self.data['Close'], self.data['Volume'])
        self.data['OBV_MA20'] = talib.MA(self.data['OBV'], timeperiod=20)
        self.data['OBV_Ratio'] = self.data['OBV'] / self.data['OBV_MA20'] - 1
        
        # 填充NaN值
        self.data = self.data.fillna(method='bfill')
        
        # 添加宏观经济数据占位符 (实际应用中需要从外部API获取)
        self.data['Economic_Indicator'] = 0  # 这里需要实际数据
        
        # 计算波动率
        self.data['Volatility'] = self.data['Close'].pct_change().rolling(20).std()
        
        # 计算相关性指标 (实际应用中需要多品种数据)
        self.data['Correlation'] = 0  # 这里需要实际数据
    
    def identify_market_regime(self):
        """识别市场状态"""
        # 使用ADX和ATR判断市场状态
        self.data['Market_Regime'] = 'Neutral'  # 默认中性市场
        
        # 强趋势市场条件
        strong_trend_condition = (self.data['ADX'] > 30) & (self.data['ATR_Pct'] > self.data['ATR_Pct'].rolling(50).mean())
        self.data.loc[strong_trend_condition, 'Market_Regime'] = 'Strong_Trend'
        
        # 震荡市场条件
        choppy_condition = (self.data['ADX'] < 20) | (self.data['ATR_Pct'] < self.data['ATR_Pct'].rolling(50).mean() * 0.7)
        self.data.loc[choppy_condition, 'Market_Regime'] = 'Choppy'
        
        # 上升趋势市场条件
        uptrend_condition = (self.data['MA20'] > self.data['MA50']) & (self.data['MA50'] > self.data['MA120'])
        self.data.loc[uptrend_condition & (self.data['Market_Regime'] != 'Strong_Trend'), 'Market_Regime'] = 'Uptrend'
        
        # 下降趋势市场条件
        downtrend_condition = (self.data['MA20'] < self.data['MA50']) & (self.data['MA50'] < self.data['MA120'])
        self.data.loc[downtrend_condition & (self.data['Market_Regime'] != 'Strong_Trend'), 'Market_Regime'] = 'Downtrend'
        
        # 使用机器学习模型增强市场状态识别
        self.enhance_market_regime_identification()
    
    def enhance_market_regime_identification(self):
        """使用简单线性模型增强市场状态识别"""
        # 准备特征
        features = ['ADX', 'ATR_Pct', 'MA20', 'MA50', 'MA120', 'Volatility']
        X = self.data[features].dropna()
        
        if len(X) > 50:  # 确保有足够的数据
            # 简单趋势预测
            y_trend = (self.data['Close'].shift(-5) / self.data['Close'] - 1).dropna()
            
            # 确保X和y长度一致
            common_index = X.index.intersection(y_trend.index)
            X = X.loc[common_index]
            y_trend = y_trend.loc[common_index]
            
            if len(X) > 10 and len(y_trend) > 10:
                # 训练简单线性模型
                model = LinearRegression()
                model.fit(X, y_trend)
                
                # 预测趋势强度
                trend_strength = model.predict(X)
                
                # 根据预测结果调整市场状态
                strong_up_trend = trend_strength > 0.02  # 预测未来5天上涨超过2%
                strong_down_trend = trend_strength < -0.02  # 预测未来5天下跌超过2%
                
                # 更新市场状态
                Strong_Trend_count = 0
                Strong_Downtrend_count = 0
                for i, idx in enumerate(X.index):
                    if strong_up_trend[i] and self.data.loc[idx, 'Market_Regime'] in ['Uptrend', 'Neutral']:
                        self.data.loc[idx, 'Market_Regime'] = 'Strong_Trend'
                        Strong_Trend_count = Strong_Trend_count + 1
                    elif strong_down_trend[i] and self.data.loc[idx, 'Market_Regime'] in ['Downtrend', 'Neutral']:
                        self.data.loc[idx, 'Market_Regime'] = 'Strong_Downtrend'
                        Strong_Downtrend_count = Strong_Downtrend_count + 1
                print(f"update {Strong_Trend_count} Strong_Trend by enhance_market_regime_identification")
                print(f"update {Strong_Downtrend_count} Strong_Downtrend by enhance_market_regime_identification")
    
    def generate_signals(self):
        """生成交易信号 - 根据市场状态自适应调整"""
        signals = pd.DataFrame(index=self.data.index)
        signals['Price'] = self.data['Close']
        signals['Signal'] = 0  # 0: Hold, 1: Buy, -1: Sell
        signals['Market_Regime'] = self.data['Market_Regime']
        
        # 基础条件
        ma_condition = (self.data['MA20'] > self.data['MA50']) & (self.data['MA50'] > self.data['MA120'])
        macd_condition = (self.data['MACD'] > self.data['MACD_Signal']) & (self.data['MACD_Hist'] > 0)
        rsi_condition = (self.data['RSI'] > 40) & (self.data['RSI'] < 70)
        bb_condition = self.data['Close'] > self.data['BB_Middle']
        volume_condition = self.data['Volume'] > self.data['Volume_MA20'] * 1.2
        cci_condition = (self.data['CCI'] > -100) & (self.data['CCI'] < 100)
        adx_condition = self.data['ADX'] > 25
        
        # 新增因子条件
        obv_condition = self.data['OBV'] > self.data['OBV_MA20']  # OBV在均线上方
        momentum_condition = self.data['Momentum'] > 0  # 正动量
        sentiment_condition = (self.data['Price_MA20_Ratio'] > -0.05) & (self.data['Price_MA50_Ratio'] > -0.08)  # 价格在均线附近或上方
        
        # 根据不同市场状态调整策略参数
        for i in range(len(signals)):
            current_regime = signals['Market_Regime'].iloc[i]
            
            # 强趋势市场 - 放宽买入条件，收紧卖出条件
            if current_regime == 'Strong_Trend':
                buy_conditions = [
                    ma_condition.iloc[i],
                    macd_condition.iloc[i] or self.data['MACD_Hist'].iloc[i] > 0,
                    rsi_condition.iloc[i] or (self.data['RSI'].iloc[i] > 30 and self.data['RSI'].iloc[i] < 80),
                    bb_condition.iloc[i],
                    volume_condition.iloc[i] or self.data['Volume_Ratio'].iloc[i] > 0.8,
                    cci_condition.iloc[i] or (self.data['CCI'].iloc[i] > -150 and self.data['CCI'].iloc[i] < 150),
                    adx_condition.iloc[i] or self.data['ADX'].iloc[i] > 20,
                    obv_condition.iloc[i],
                    momentum_condition.iloc[i],
                    sentiment_condition.iloc[i]
                ]
                
                # 强趋势市场中，收紧卖出条件
                sell_condition = (
                    (self.data['MA20'].iloc[i] < self.data['MA50'].iloc[i] * 0.98) or  # MA20明显下穿MA50
                    (self.data['MACD'].iloc[i] < self.data['MACD_Signal'].iloc[i] * 0.9) or  # MACD明显死叉
                    (self.data['RSI'].iloc[i] > 80 or self.data['RSI'].iloc[i] < 30) or  # RSI极端值
                    (self.data['Close'].iloc[i] < self.data['BB_Lower'].iloc[i] * 0.98) or  # 价格明显跌破布林带下轨
                    (self.data['ADX'].iloc[i] < 15)  # 趋势明显减弱
                )
                
                # 强趋势市场中，满足5个条件即可买入
                if sum(buy_conditions) >= 5:
                    signals.iloc[i, signals.columns.get_loc('Signal')] = 1
                elif sell_condition:
                    signals.iloc[i, signals.columns.get_loc('Signal')] = -1
            
            # 震荡市场 - 收紧买入条件，放宽卖出条件
            elif current_regime == 'Choppy':
                buy_conditions = [
                    ma_condition.iloc[i],
                    macd_condition.iloc[i],
                    rsi_condition.iloc[i] and self.data['RSI'].iloc[i] < 60,
                    bb_condition.iloc[i] and self.data['BB_Position'].iloc[i] > 0.5,
                    volume_condition.iloc[i],
                    cci_condition.iloc[i],
                    adx_condition.iloc[i],
                    obv_condition.iloc[i],
                    momentum_condition.iloc[i],
                    sentiment_condition.iloc[i]
                ]
                
                # 震荡市场中，需要满足7个条件才买入
                if sum(buy_conditions) >= 7:
                    signals.iloc[i, signals.columns.get_loc('Signal')] = 1
                
                # 震荡市场中，放宽卖出条件
                sell_condition = (
                    (self.data['MA20'].iloc[i] < self.data['MA50'].iloc[i]) or
                    (self.data['MACD'].iloc[i] < self.data['MACD_Signal'].iloc[i]) or
                    (self.data['RSI'].iloc[i] > 70 or self.data['RSI'].iloc[i] < 40) or
                    (self.data['Close'].iloc[i] < self.data['BB_Lower'].iloc[i]) or
                    (self.data['ADX'].iloc[i] < 20)
                )
                
                if sell_condition:
                    signals.iloc[i, signals.columns.get_loc('Signal')] = -1
                elif sum(buy_conditions) >= 7:
                    signals.iloc[i, signals.columns.get_loc('Signal')] = 1
            
            # 上升趋势市场 - 适度条件
            elif current_regime == 'Uptrend':
                buy_conditions = [
                    ma_condition.iloc[i],
                    macd_condition.iloc[i],
                    rsi_condition.iloc[i],
                    bb_condition.iloc[i],
                    volume_condition.iloc[i],
                    cci_condition.iloc[i],
                    adx_condition.iloc[i],
                    obv_condition.iloc[i],
                    momentum_condition.iloc[i],
                    sentiment_condition.iloc[i]
                ]
                
                # 适度卖出条件
                sell_condition = (
                    (self.data['MA20'].iloc[i] < self.data['MA50'].iloc[i]) or
                    (self.data['MACD'].iloc[i] < self.data['MACD_Signal'].iloc[i]) or
                    (self.data['RSI'].iloc[i] > 75 or self.data['RSI'].iloc[i] < 35) or
                    (self.data['Close'].iloc[i] < self.data['BB_Lower'].iloc[i]) or
                    (self.data['ADX'].iloc[i] < 18)
                )
                
                
                # 上升趋势中，满足6个条件即可买入
                if sum(buy_conditions) >= 6:
                    signals.iloc[i, signals.columns.get_loc('Signal')] = 1
                elif sell_condition:
                    signals.iloc[i, signals.columns.get_loc('Signal')] = -1
            
            # 下降趋势市场 - 谨慎操作，主要做空或观望
            elif current_regime in ['Downtrend', 'Strong_Downtrend']:
                # 在下降趋势中，只考虑超跌反弹机会
                oversold_condition = (
                    (self.data['RSI'].iloc[i] < 30) and
                    (self.data['Close'].iloc[i] < self.data['BB_Lower'].iloc[i] * 1.02) and
                    (self.data['CCI'].iloc[i] < -100) and
                    (self.data['Volume'].iloc[i] > self.data['Volume_MA20'].iloc[i] * 1.5) and
                    (self.data['OBV_Ratio'].iloc[i] > -0.1)  # OBV没有大幅下降
                )
                
                # 下降趋势中，任何反弹都是卖出机会
                rebound_condition = (
                    (self.data['RSI'].iloc[i] > 50) or
                    (self.data['Close'].iloc[i] > self.data['MA20'].iloc[i]) or
                    (self.data['MACD_Hist'].iloc[i] < 0) or
                    (self.data['OBV_Ratio'].iloc[i] < -0.05)  # OBV开始下降
                )
                
                if rebound_condition and (i == 0 or signals['Signal'].iloc[i-1] == 1):
                    signals.iloc[i, signals.columns.get_loc('Signal')] = -1
                elif oversold_condition:
                    signals.iloc[i, signals.columns.get_loc('Signal')] = 1
        
        # 确保信号连续性
        #signals['Signal'] = signals['Signal'].replace(to_replace=0, method='ffill')
        #signals.to_csv('data.csv', index=False, encoding='utf-8-sig') 
        
        return signals
    
    def calculate_step_size(self, signals, risk_per_trade=0.01, max_step= 2):
        """每次买入卖出的步进大小 - 基于波动率和市场状态"""
        signals['Step'] = 0.0
        
        # 计算基于波动率的仓位
        atr_normalized = self.data['ATR_Pct'].rolling(20).mean()
        volatility_factor = 0.02 / atr_normalized  # 目标波动率2%
        volatility_factor = volatility_factor.clip(upper=2.0, lower=0.5)  # 限制范围
        
        # 根据市场状态调整仓位
        for i in range(len(signals)):
            regime = signals['Market_Regime'].iloc[i]
            if signals['Signal'].iloc[i] == 1:  # 买入信号
                
                if regime == 'Strong_Trend':
                    step_size = 0.8 * volatility_factor.iloc[i]
                elif regime == 'Uptrend':
                    step_size = 0.6 * volatility_factor.iloc[i]
                elif regime == 'Choppy':
                    step_size = 0.3 * volatility_factor.iloc[i]
                elif regime == 'Downtrend':
                    step_size = 0.2 * volatility_factor.iloc[i]
                elif regime == 'Strong_Downtrend':
                    step_size = 0.1 * volatility_factor.iloc[i]
                else:
                    step_size = 0.4 * volatility_factor.iloc[i]
                
                # 考虑相关性调整 (这里简化处理)
                correlation_adjustment = 1.0  # 默认无调整
                if abs(self.data['Correlation'].iloc[i]) > 0.7:  # 高相关性
                    correlation_adjustment = 0.7  # 降低仓位
                
                signals.iloc[i, signals.columns.get_loc('Step')] = min(
                    step_size * correlation_adjustment, max_step)
            elif signals['Signal'].iloc[i] == -1: # 卖出信号
                if regime == 'Strong_Trend':
                    step_size = 0.1 * volatility_factor.iloc[i]
                elif regime == 'Uptrend':
                    step_size = 0.2 * volatility_factor.iloc[i]
                elif regime == 'Choppy':
                    step_size = 0.3 * volatility_factor.iloc[i]
                elif regime == 'Downtrend':
                    step_size = 0.6 * volatility_factor.iloc[i]
                elif regime == 'Strong_Downtrend':
                    step_size = 0.8 * volatility_factor.iloc[i]
                else:
                    step_size = 0.4 * volatility_factor.iloc[i]
                signals.iloc[i, signals.columns.get_loc('Step')] = step_size
            else:
                signals.iloc[i, signals.columns.get_loc('Step')] = 0
                
        signals.to_csv('pos.csv', index=True, encoding='utf-8-sig')
        
        return signals
    
    def backtest(self, signals, initial_capital=100000.0, transaction_cost=0.005, 
                max_daily_loss=0.05, max_drawdown=0.2):
        """回测策略 - 加入动态仓位管理和风险控制"""
        # 初始化
        portfolio = pd.DataFrame(index=signals.index)
        portfolio['Price'] = signals['Price']
        portfolio['Signal'] = signals['Signal']
        portfolio['Step'] = signals['Step']
        portfolio['Holdings'] = 0.0 # 持仓
        portfolio['Cash'] = initial_capital
        portfolio['Total'] = initial_capital
        portfolio['Returns'] = 0.0
        portfolio['Peak'] = initial_capital  # 资产峰值
        portfolio['Drawdown'] = 0.0  # 回撤
        
        # 止损相关变量
        portfolio['Stop_Loss_Price'] = 0.0  # 止损价格
        portfolio['Time_In_Trade'] = 0  # 持仓天数
        portfolio['Entry_Price'] = 0.0  # 持仓成本价格
        
        # 风险控制标志
        trading_halted = False
        halt_reason = ""

        #记录总成本
        total_cost = 0
        
        # 记录交易
        trades = pd.DataFrame(columns=['Date', 'Type', 'Price', 'Shares', 'Value', 'Cash', 'Reason'])
        #a = 1
        portfolio = portfolio.fillna(0)

        for i in range(2, len(portfolio)):
            current_date = portfolio.index[i]
            prev_date = portfolio.index[i-1]
            #prev_pre_date = portfolio.index[i-2]
            current_price = portfolio['Price'].iloc[i]
            current_signal = portfolio['Signal'].iloc[i]
            current_step = portfolio['Step'].iloc[i] / 10 # 原step是每1个单位代表购买或者卖出10% 所以这里乘以10
            
            # 继承前一天的持仓和现金
            portfolio.loc[current_date, 'Holdings'] = portfolio.loc[prev_date, 'Holdings']
            portfolio.loc[current_date, 'Cash'] = portfolio.loc[prev_date, 'Cash']
            portfolio.loc[current_date, 'Time_In_Trade'] = portfolio.loc[prev_date, 'Time_In_Trade']
            portfolio.loc[current_date, 'Entry_Price'] = portfolio.loc[prev_date, 'Entry_Price']
            portfolio.loc[current_date, 'Stop_Loss_Price'] = portfolio.loc[prev_date, 'Stop_Loss_Price']
            
            # 更新总资产和收益
            portfolio.loc[current_date, 'Total'] = portfolio.loc[current_date, 'Cash'] + \
                                                portfolio.loc[current_date, 'Holdings'] * current_price
            portfolio.loc[current_date, 'Returns'] = (portfolio.loc[current_date, 'Total'] / 
                                                portfolio.loc[prev_date, 'Total'] - 1)

            
            # 更新资产峰值和回撤
            portfolio.loc[current_date, 'Peak'] = max(portfolio.loc[prev_date, 'Peak'], 
                                                    portfolio.loc[current_date, 'Total'])
            portfolio.loc[current_date, 'Drawdown'] = (portfolio.loc[current_date, 'Total'] - 
                                                     portfolio.loc[current_date, 'Peak']) / portfolio.loc[current_date, 'Peak']

            #用0 替代空值
            current_step = 0 if current_step is None else current_step
            #if a:
            #    print(f"data::::{current_date}   {portfolio.loc[current_date, 'Holdings']}")
            #当step 小于等于0时不用操作
            if current_step <= 0:
                continue
            
            # 检查风险控制
            if trading_halted:
                # 如果交易被暂停， 检查是否为趋势行情，是则更新暂停，改为继续操作，只更新资产价值
                current_regime = signals['Market_Regime'].iloc[i]
                if current_regime in ['Uptrend', 'Strong_Trend']:
                    trading_halted = False
                else:
                    #portfolio.loc[current_date, 'Total'] = portfolio.loc[current_date, 'Cash'] + \
                    #                                    portfolio.loc[current_date, 'Holdings'] * current_price

                    continue

            #### 备忘： 后面每次清仓 资产峰值清0，持仓天数清0. 另外，买入价格改为成本价格
            # 检查最大回撤
            if portfolio.loc[current_date, 'Drawdown'] < -max_drawdown:
                trading_halted = True
                halt_reason = f"Max Drawdown exceeded: {portfolio.loc[current_date, 'Drawdown']*100:.2f}%"
                # 清空所有持仓
                if portfolio.loc[current_date, 'Holdings'] > 0:
                    shares_to_sell = portfolio.loc[current_date, 'Holdings']
                    sale_proceeds = shares_to_sell * current_price * (1 - transaction_cost)
                    portfolio.loc[current_date, 'Cash'] = portfolio.loc[current_date, 'Cash'] + sale_proceeds
                    portfolio.loc[current_date, 'Holdings'] = 0
                    
                    # 记录交易
                    trades = trades._append({
                        'Date': current_date,
                        'Type': 'CLEAR',
                        'Price': current_price,
                        'Shares': shares_to_sell,
                        'Value': sale_proceeds,
                        'Cash': portfolio.loc[current_date, 'Cash'],
                        'Reason': halt_reason
                    }, ignore_index=True)
                
                total_cost = 0
                portfolio.loc[current_date, 'Peak'] = 0
                portfolio.loc[current_date, 'Entry_Price'] = 0
                portfolio.loc[current_date, 'Time_In_Trade'] = 0
                
                portfolio.loc[current_date, 'Total'] = portfolio.loc[current_date, 'Cash'] + \
                                                     portfolio.loc[current_date, 'Holdings'] * current_price
                portfolio.loc[current_date, 'Returns'] = (portfolio.loc[current_date, 'Total'] / 
                                                        portfolio.loc[prev_date, 'Total'] - 1)
                continue
            
            # 检查每日损失限额
            daily_return = (portfolio.loc[current_date, 'Total'] / portfolio.loc[prev_date, 'Total'] - 1)
            if daily_return < -max_daily_loss:
                trading_halted = True
                halt_reason = f"Daily Loss Limit exceeded: {daily_return*100:.2f}%"
                # 清空所有持仓
                if portfolio.loc[current_date, 'Holdings'] > 0:
                    shares_to_sell = portfolio.loc[current_date, 'Holdings']
                    sale_proceeds = shares_to_sell * current_price * (1 - transaction_cost)
                    portfolio.loc[current_date, 'Cash'] = portfolio.loc[current_date, 'Cash'] + sale_proceeds
                    portfolio.loc[current_date, 'Holdings'] = 0
                    
                    # 记录交易
                    trades = trades._append({
                        'Date': current_date,
                        'Type': 'CLEAR1',
                        'Price': current_price,
                        'Shares': shares_to_sell,
                        'Value': sale_proceeds,
                        'Cash': portfolio.loc[current_date, 'Cash'],
                        'Reason': halt_reason
                    }, ignore_index=True)
                
                total_cost = 0
                portfolio.loc[current_date, 'Peak'] = 0
                portfolio.loc[current_date, 'Entry_Price'] = 0
                portfolio.loc[current_date, 'Time_In_Trade'] = 0
                    
                portfolio.loc[current_date, 'Total'] = portfolio.loc[current_date, 'Cash'] + \
                                                     portfolio.loc[current_date, 'Holdings'] * current_price
                portfolio.loc[current_date, 'Returns'] = (portfolio.loc[current_date, 'Total'] / 
                                                        portfolio.loc[prev_date, 'Total'] - 1)
                continue
            
            # 买入信号
            if current_signal == 1:
                
                # 计算目标市值
                target_value = initial_capital * current_step
                target_value = min(target_value, portfolio.loc[current_date, 'Cash'])
                # 计算需要购买的股数
                shares_to_buy = target_value / current_price
                shares_to_buy = round(shares_to_buy, 2)
                cost = shares_to_buy * current_price * (1 + transaction_cost)
                
                total_cost = total_cost + cost
                
                # 更新持仓和现金
                portfolio.loc[current_date, 'Holdings'] = portfolio.loc[current_date, 'Holdings'] + shares_to_buy
                portfolio.loc[current_date, 'Cash'] = portfolio.loc[current_date, 'Cash'] - cost
                portfolio.loc[current_date, 'Entry_Price'] = total_cost / portfolio.loc[current_date, 'Holdings']
                portfolio.loc[current_date, 'Time_In_Trade'] = portfolio.loc[current_date, 'Time_In_Trade'] + 1
                    
                portfolio.loc[current_date, 'Total'] = portfolio.loc[current_date, 'Cash'] + \
                                                     portfolio.loc[current_date, 'Holdings'] * current_price
                portfolio.loc[current_date, 'Returns'] = (portfolio.loc[current_date, 'Total'] / 
                                                        portfolio.loc[prev_date, 'Total'] - 1)
                # 设置基于ATR的止损（暂时没有用到）
                # atr_stop = 2 * self.data['ATR'].iloc[i]  # 2倍ATR止损
                # portfolio.loc[current_date, 'Stop_Loss_Price'] = current_price - atr_stop
                
                # 记录交易
                trades = trades._append({
                    'Date': current_date,
                    'Type': 'BUY',
                    'Price': current_price,
                    'Shares': shares_to_buy,
                    'Value': cost,
                    'Cash': portfolio.loc[current_date, 'Cash'],
                    'Reason': 'Buy Signal'
                }, ignore_index=True)
            
            # 卖出信号
            elif current_signal == -1:
                # 计算目标市值
                
                target_value = initial_capital * current_step
                target_value = min(target_value, portfolio.loc[current_date, 'Holdings'] * current_price)
                # 计算需要购买的股数
                shares_to_sell = target_value / current_price
                shares_to_sell = round(shares_to_sell, 2)
                sell = shares_to_sell * current_price * (1 - transaction_cost)
                
                #total_cost = total_cost * (1 - sell/(portfolio.loc[current_date, 'Holdings'] * current_price))
                total_cost = total_cost - shares_to_sell * portfolio.loc[current_date, 'Entry_Price']
                # 更新持仓和现金
                portfolio.loc[current_date, 'Holdings'] = portfolio.loc[current_date, 'Holdings'] - shares_to_sell
                portfolio.loc[current_date, 'Cash'] = portfolio.loc[current_date, 'Cash'] + sell

                if portfolio.loc[current_date, 'Holdings'] > 0:
                    portfolio.loc[current_date, 'Entry_Price'] = total_cost / portfolio.loc[current_date, 'Holdings']
                    portfolio.loc[current_date, 'Time_In_Trade'] = portfolio.loc[current_date, 'Time_In_Trade'] + 1
                else:
                    total_cost = 0
                    portfolio.loc[current_date, 'Peak'] = 0
                    portfolio.loc[current_date, 'Entry_Price'] = 0
                    portfolio.loc[current_date, 'Time_In_Trade'] = 0

                portfolio.loc[current_date, 'Total'] = portfolio.loc[current_date, 'Cash'] + \
                                                     portfolio.loc[current_date, 'Holdings'] * current_price
                portfolio.loc[current_date, 'Returns'] = (portfolio.loc[current_date, 'Total'] / 
                                                        portfolio.loc[prev_date, 'Total'] - 1)
                
                # 设置基于ATR的止损（暂时没有用）
                # atr_stop = 2 * self.data['ATR'].iloc[i]  # 2倍ATR止损
                # portfolio.loc[current_date, 'Stop_Loss_Price'] = current_price - atr_stop


                # 卖出所有持仓
                #shares_to_sell = portfolio.loc[current_date, 'Holdings']
                #sale_proceeds = shares_to_sell * current_price * (1 - transaction_cost)
                
                # 更新持仓和现金
                #portfolio.loc[current_date, 'Holdings'] = 0
                #portfolio.loc[current_date, 'Cash'] = portfolio.loc[current_date, 'Cash'] + sale_proceeds
                #portfolio.loc[current_date, 'Time_In_Trade'] = 0
                #portfolio.loc[current_date, 'Entry_Price'] = 0.0
                #portfolio.loc[current_date, 'Stop_Loss_Price'] = 0.0
                
                # 记录交易
                trades = trades._append({
                    'Date': current_date,
                    'Type': 'SELL',
                    'Price': current_price,
                    'Shares': shares_to_sell,
                    'Value': sell,
                    'Cash': portfolio.loc[current_date, 'Cash'],
                    'Reason': 'Sell Signal'
                }, ignore_index=True)
            '''
            # 检查止损条件
            else:  # 当前持有仓位
                # 更新持仓天数
                if portfolio.loc[current_date, 'Holdings'] > 0:
                    portfolio.loc[current_date, 'Time_In_Trade'] = portfolio.loc[current_date, 'Time_In_Trade'] + 1
                else:
                    continue
                
                # 基于ATR的动态止损
                current_stop_loss = portfolio.loc[current_date, 'Stop_Loss_Price']
                if current_price < current_stop_loss:
                    # 触发止损
                    shares_to_sell = portfolio.loc[current_date, 'Holdings']
                    sale_proceeds = shares_to_sell * current_price * (1 - transaction_cost)
                    
                    # 更新持仓和现金
                    portfolio.loc[current_date, 'Holdings'] = 0
                    portfolio.loc[current_date, 'Cash'] = portfolio.loc[current_date, 'Cash'] + sale_proceeds
                    portfolio.loc[current_date, 'Time_In_Trade'] = 0
                    portfolio.loc[current_date, 'Entry_Price'] = 0.0
                    portfolio.loc[current_date, 'Stop_Loss_Price'] = 0.0
                    
                    # 记录交易
                    trades = trades._append({
                        'Date': current_date,
                        'Type': 'SELL',
                        'Price': current_price,
                        'Shares': shares_to_sell,
                        'Value': sale_proceeds,
                        'Cash': portfolio.loc[current_date, 'Cash'],
                        'Reason': 'ATR Stop Loss'
                    }, ignore_index=True)
                
                # 时间止损 (持仓超过20天且无盈利)
                elif (portfolio.loc[current_date, 'Time_In_Trade'] > 20 and 
                      current_price < portfolio.loc[current_date, 'Entry_Price']):
                    # 触发时间止损
                    shares_to_sell = portfolio.loc[current_date, 'Holdings']
                    sale_proceeds = shares_to_sell * current_price * (1 - transaction_cost)
                    
                    # 更新持仓和现金
                    portfolio.loc[current_date, 'Holdings'] = 0
                    portfolio.loc[current_date, 'Cash'] = portfolio.loc[current_date, 'Cash'] + sale_proceeds
                    portfolio.loc[current_date, 'Time_In_Trade'] = 0
                    portfolio.loc[current_date, 'Entry_Price'] = 0.0
                    portfolio.loc[current_date, 'Stop_Loss_Price'] = 0.0
                    
                    # 记录交易
                    trades = trades._append({
                        'Date': current_date,
                        'Type': 'SELL',
                        'Price': current_price,
                        'Shares': shares_to_sell,
                        'Value': sale_proceeds,
                        'Cash': portfolio.loc[current_date, 'Cash'],
                        'Reason': 'Time Stop Loss'
                    }, ignore_index=True)
                
                # 更新动态止损 (追踪止损)
                else:
                    # 基于最高价的回撤止损
                    current_peak = max(portfolio.loc[current_date, 'Entry_Price'], 
                                      portfolio.loc[prev_date, 'Stop_Loss_Price'] + 0.5 * self.data['ATR'].iloc[i])
                    new_stop_loss = current_peak - 1.5 * self.data['ATR'].iloc[i]  # 1.5倍ATR追踪止损
                    
                    # 只向上移动止损
                    if new_stop_loss > portfolio.loc[current_date, 'Stop_Loss_Price']:
                        portfolio.loc[current_date, 'Stop_Loss_Price'] = new_stop_loss
            '''
            # 更新总资产和收益
            portfolio.loc[current_date, 'Total'] = portfolio.loc[current_date, 'Cash'] + \
                                                 portfolio.loc[current_date, 'Holdings'] * current_price
            portfolio.loc[current_date, 'Returns'] = (portfolio.loc[current_date, 'Total'] / 
                                                    portfolio.loc[prev_date, 'Total'] - 1)
        
        # 如果交易被暂停，添加原因到portfolio 暂停只是暂时的，并不是今后再不交易
        # if trading_halted:
        #    portfolio['Trading_Halted'] = True
        #    portfolio['Halt_Reason'] = halt_reason
        #else:
        #    portfolio['Trading_Halted'] = False
        #    portfolio['Halt_Reason'] = ""

        trades.to_csv('trades.csv', index=False, encoding='utf-8-sig') 
        portfolio.to_csv('portolio.csv', index=True, encoding='utf-8-sig')
        return portfolio, trades
    
    def walk_forward_optimization(self, initial_period=252, test_period=63, steps=5):
        """Walk-Forward参数优化"""
        results = []
        
        for step in range(steps):
            # 划分训练集和测试集
            train_end = initial_period + step * test_period
            test_start = train_end
            test_end = test_start + test_period
            
            if test_end > len(self.data):
                break
            
            # 训练集数据
            train_data = self.data.iloc[:train_end]
            
            # 这里可以添加参数优化逻辑
            # 例如：优化不同市场状态下的买入条件阈值
            
            # 测试集数据
            test_data = self.data.iloc[test_start:test_end]
            
            # 在测试集上回测
            # 这里简化处理，实际应该使用优化后的参数生成信号并回测
            
            # 记录结果
            result = {
                'Step': step,
                'Train_Period': f"{self.data.index[0].date()} to {self.data.index[train_end-1].date()}",
                'Test_Period': f"{self.data.index[test_start].date()} to {self.data.index[test_end-1].date()}",
                # 这里添加性能指标
            }
            
            results.append(result)
        
        return pd.DataFrame(results)
    
    def calculate_performance_metrics(self, portfolio):
        """计算绩效指标"""
        # 总收益
        total_return = (portfolio['Total'].iloc[-1] / portfolio['Total'].iloc[0] - 1) * 100
        
        # 年化收益
        days = (portfolio.index[-1] - portfolio.index[0]).days
        if days > 0:
            annualized_return = (1 + total_return/100) ** (365/days) - 1
        else:
            annualized_return = 0
        
        # 最大回撤
        peak = portfolio['Total'].cummax()
        drawdown = (portfolio['Total'] - peak) / peak
        max_drawdown = drawdown.min() * 100
        
        # 夏普比率 (假设无风险利率3%)
        excess_returns = portfolio['Returns'] - 0.03/252
        if excess_returns.std() > 0:
            sharpe_ratio = np.sqrt(252) * excess_returns.mean() / excess_returns.std()
        else:
            sharpe_ratio = 0
        
        # 胜率
        trade_returns = portfolio['Returns'][portfolio['Signal'].diff() != 0]
        if len(trade_returns) > 0:
            win_rate = (trade_returns > 0).mean() * 100
        else:
            win_rate = 0
        
        # 盈亏比
        winning_trades = trade_returns[trade_returns > 0]
        losing_trades = trade_returns[trade_returns < 0]
        if len(losing_trades) > 0 and len(winning_trades) > 0:
            profit_loss_ratio = winning_trades.mean() / abs(losing_trades.mean())
        else:
            profit_loss_ratio = 0
        
        # 索提诺比率 (只考虑下行风险)
        downside_returns = excess_returns[excess_returns < 0]
        if downside_returns.std() > 0:
            sortino_ratio = np.sqrt(252) * excess_returns.mean() / downside_returns.std()
        else:
            sortino_ratio = 0
        
        return {
            'Total Return (%)': total_return,
            'Annualized Return (%)': annualized_return * 100,
            'Max Drawdown (%)': max_drawdown,
            'Sharpe Ratio': sharpe_ratio,
            'Sortino Ratio': sortino_ratio,
            'Win Rate (%)': win_rate,
            'Profit/Loss Ratio': profit_loss_ratio,
            'Number of Trades': len(trade_returns)
        }
    
    def plot_results(self, portfolio, signals):
        """绘制结果"""
        fig, axes = plt.subplots(3, 1, figsize=(16, 16))
        
        # 资产曲线
        axes[0].plot(portfolio.index, portfolio['Total'], label='Strategy NAV', linewidth=2)
        axes[0].plot(portfolio.index, portfolio['Total'].iloc[0] * 
                    (1 + self.data['Close'].pct_change().cumsum().fillna(0)), label='Benchmark (CSI300)', alpha=0.7)
        
        # 标记买卖点
        buy_signals = signals[signals['Signal'] == 1]
        sell_signals = signals[signals['Signal'] == -1]
        axes[0].scatter(buy_signals.index, portfolio.loc[buy_signals.index, 'Total'], 
                       color='green', marker='^', s=100, label='Buy', alpha=0.7)
        axes[0].scatter(sell_signals.index, portfolio.loc[sell_signals.index, 'Total'], 
                       color='red', marker='v', s=100, label='Sell', alpha=0.7)
        
        axes[0].set_title('Strategy Performance')
        axes[0].set_ylabel('Portfolio Value')
        axes[0].legend()
        axes[0].grid(True)
        
        # 市场状态
        regime_colors = {
            'Strong_Trend': 'green',
            'Uptrend': 'lightgreen',
            'Neutral': 'gray',
            'Choppy': 'orange',
            'Downtrend': 'red',
            'Strong_Downtrend': 'darkred'
        }
        
        for regime, color in regime_colors.items():
            regime_data = signals[signals['Market_Regime'] == regime]
            if not regime_data.empty:
                axes[1].plot(regime_data.index, self.data.loc[regime_data.index, 'Close'], 
                            color=color, label=regime, linestyle='', marker='o', markersize=2)
        
        axes[1].set_title('Market Regimes')
        axes[1].set_ylabel('Price')
        axes[1].legend()
        axes[1].grid(True)

        # 持有收益率
        axes[2].plot(portfolio.index, portfolio['Entry_Price']/portfolio['Price'] -1 , label='holding period return(%)', linewidth=2)
        axes[2].set_title('holding period return')
        axes[2].set_ylabel('Return (%)')
        axes[2].grid(True)
        '''
        # 仓位变化
        axes[2].plot(portfolio.index, portfolio['Step'] * 100, label='Position (%)', linewidth=2)
        axes[2].set_title('Position Sizing')
        axes[2].set_ylabel('Position (%)')
        axes[2].grid(True)
        
        # 回撤曲线
        axes[3].plot(portfolio.index, portfolio['Drawdown'] * 100, label='Drawdown (%)', color='red', linewidth=2)
        axes[3].axhline(y=0, color='black', linestyle='-', alpha=0.3)
        axes[3].set_title('Drawdown')
        axes[3].set_ylabel('Drawdown (%)')
        axes[3].set_xlabel('Date')
        axes[3].grid(True)
        '''
        plt.tight_layout()
        plt.show()

# 使用示例
if __name__ == "__main__":
    # 创建策略实例
    strategy = EnhancedCSI300Strategy(symbol="sh000300", start_date="20050101")
    
    # 生成交易信号
    signals = strategy.generate_signals()
    
    # 计算每次操作的步进值大小
    signals = strategy.calculate_step_size(signals)
    
    # 回测策略
    portfolio, trades = strategy.backtest(signals)
    
    # 计算绩效指标
    metrics = strategy.calculate_performance_metrics(portfolio)
    print("策略绩效指标:")
    for key, value in metrics.items():
        if '%' in key:
            print(f"{key}: {value:.2f}")
        else:
            print(f"{key}: {value:.4f}")
    
    # 绘制结果
    strategy.plot_results(portfolio, signals)
    
    # 显示最近交易信号
    print("\n最近交易信号:")
    print("saved")
    
    # 显示交易历史
    print("\n交易历史:")
    print("saved")
    
    # 执行Walk-Forward优化
    # wfo_results = strategy.walk_forward_optimization()
    # print("\nWalk-Forward优化结果:")
    # print(wfo_results)